# 导入必要的模块
from PyQt6.QtWidgets import (QStyledItemDelegate, QSpinBox, QComboBox, 
                            QMainWindow, QTableView, QApplication)
from PyQt6.QtCore import Qt, QAbstractTableModel
import sys
import warnings
from PyQt6 import sip

# 自定义表格模型类，继承自QAbstractTableModel
class CustomTableModel(QAbstractTableModel):
    def __init__(self, data, headers):
        """
        初始化自定义表格模型
        :param data: 二维列表，包含表格数据
        :param headers: 列表，包含列标题
        """
        super().__init__()
        self._data = data      # 存储表格数据
        self._headers = headers  # 存储列标题
    
    def rowCount(self, parent=None):
        """返回表格的行数"""
        return len(self._data)
    
    def columnCount(self, parent=None):
        """返回表格的列数"""
        return len(self._headers)
    
    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
        """
        返回指定索引的数据
        :param index: QModelIndex对象，包含行列信息
        :param role: 数据角色，决定数据的用途(显示、编辑、装饰等)
        """
        if not index.isValid():  # 检查索引是否有效
            return None
        
        if role in (Qt.ItemDataRole.DisplayRole, Qt.ItemDataRole.EditRole):  # 处理显示和编辑角色
            return self._data[index.row()][index.column()]
        
        return None
    
    def headerData(self, section, orientation, role=Qt.ItemDataRole.DisplayRole):
        """
        返回表头数据
        :param section: 列号(水平表头)或行号(垂直表头)
        :param orientation: 表头方向(水平或垂直)
        :param role: 数据角色
        """
        if role == Qt.ItemDataRole.DisplayRole and orientation == Qt.Orientation.Horizontal:
            return self._headers[section]  # 返回水平表头的文本
        
        return None

    def flags(self, index):
        """返回每个单元格的标志"""
        default_flags = super().flags(index)
        if index.isValid():
            return default_flags | Qt.ItemFlag.ItemIsEditable
        return default_flags

    def setData(self, index, value, role=Qt.ItemDataRole.EditRole):
        """
        设置指定索引的数据
        :param index: QModelIndex对象
        :param value: 要设置的值
        :param role: 数据角色
        :return: 是否设置成功
        """
        if not index.isValid() or role != Qt.ItemDataRole.EditRole:
            return False
        
        # 更新数据
        self._data[index.row()][index.column()] = value
        # 发出数据改变信号，通知视图更新
        self.dataChanged.emit(index, index, [role])
        return True

# 自定义委托类，继承自QStyledItemDelegate
class AgeDelegate(QStyledItemDelegate):
    """
    自定义委托类，用于控制特定列的编辑方式
    第1列(年龄)使用QSpinBox编辑
    第3列(性别)使用QComboBox编辑
    其他列使用默认编辑方式
    """
    
    def createEditor(self, parent, option, index):
        """
        创建用于编辑的控件
        :param parent: 父控件
        :param option: 样式选项
        :param index: 模型索引
        :return: 创建的编辑器控件
        """
        if index.column() == 1:  # 如果是年龄列(第1列，从0开始)
            # 创建QSpinBox作为编辑器
            editor = QSpinBox(parent)
            editor.setMinimum(0)    # 设置最小值
            editor.setMaximum(150)  # 设置最大值
            editor.setFrame(True)   # 添加这行，确保编辑器可见
            editor.setFocusPolicy(Qt.FocusPolicy.StrongFocus)  # 确保能获取焦点
            return editor
        elif index.column() == 3:  # 如果是性别列(第3列)
            # 创建QComboBox作为编辑器
            editor = QComboBox(parent)
            editor.addItems(["男", "女", "其他"])  # 添加选项
            return editor
        else:
            # 其他列使用父类的默认编辑器(通常是QLineEdit)
            return super().createEditor(parent, option, index)
    
    def setEditorData(self, editor, index):
        """
        将模型中的数据加载到编辑器中
        :param editor: 编辑器控件
        :param index: 模型索引
        """
        value = index.model().data(index, Qt.ItemDataRole.EditRole)
        
        if index.column() == 1:  # 年龄列
            # 处理可能的None值，设置默认值0
            try:
                editor.setValue(int(value) if value is not None else 0)
            except (ValueError, TypeError):
                editor.setValue(0)
        elif index.column() == 3:  # 性别列
            # 处理可能的None值，设置默认第一个选项
            editor.setCurrentText(value if value is not None else "男")
        else:
            # 其他列使用父类的默认处理方式
            super().setEditorData(editor, index)
    
    def setModelData(self, editor, model, index):
        """
        将编辑器中的数据保存回模型
        :param editor: 编辑器控件
        :param model: 数据模型
        :param index: 模型索引
        """
        if index.column() == 1:  # 年龄列
            # 从QSpinBox获取值并保存到模型
            model.setData(index, editor.value(), Qt.ItemDataRole.EditRole)
        elif index.column() == 3:  # 性别列
            # 从QComboBox获取当前文本并保存到模型
            model.setData(index, editor.currentText(), Qt.ItemDataRole.EditRole)
        else:
            # 其他列使用父类的默认处理方式
            super().setModelData(editor, model, index)

# 主窗口类，用于演示自定义委托
class DelegateDemo(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Delegate 示例")
        self.resize(600, 400)
        
        # 准备表格数据
        data = [
            ["张三", "25", "工程师", "男"],  # 每行数据
            ["李四", "30", "设计师", "女"],
            ["王五", "28", "教师", "男"],
            ["赵六", "35", "医生", "女"]
        ]
        headers = ["姓名", "年龄", "职业", "性别"]  # 列标题
        
        # 创建自定义模型并设置数据
        model = CustomTableModel(data, headers)
        
        # 创建表格视图
        view = QTableView()
        view.setModel(model)
        view.setEditTriggers(QTableView.EditTrigger.DoubleClicked | 
                           QTableView.EditTrigger.EditKeyPressed)  # 设置编辑触发方式
        
        # 设置自定义委托
        view.setItemDelegate(AgeDelegate())
        
        # 将表格视图设置为中心部件
        self.setCentralWidget(view)

# 应用程序入口
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = DelegateDemo()
    window.show()
    sys.exit(app.exec())